package class15to16_UnionSet;

import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;

public class UnionFind<T> {
    public static class Node<T> {
        T value;

        public Node(T t) {
            value = t;
        }
    }
    private HashMap<T,Node<T>> nodeMap= new HashMap<>();
    private HashMap<Node<T>,Node<T>> leaderMap = new HashMap<>();
    private HashMap<Node<T>,Integer> sizeMap = new HashMap<>();

    public UnionFind(T[] ts){
        for(T t: ts){
            Node<T> node = new Node<T>(t);
            nodeMap.put(t,node);
            leaderMap.put(node,node);
            sizeMap.put(node,1);
        }
    }

    public Node<T> findRepresent(Node<T> n){

        LinkedList<Node<T>> stack = new LinkedList<>();
        while(leaderMap.get(n)!=n){
            stack.push(n);
            n=leaderMap.get(n);
        }
        //使结构更扁平，效率更高
        while(!stack.isEmpty()){
            leaderMap.put(stack.pop(),n);
        }
        return n;
    }

    public boolean isSameSet(T a, T b){
        Node<T> representA = findRepresent(nodeMap.get(a));
        Node<T> representB = findRepresent(nodeMap.get(b));
        return representA == representB;
    }

    public void union(T a, T b){
        Node<T> nodeA = nodeMap.get(a);
        Node<T> nodeB = nodeMap.get(b);
        Node<T> leaderA = findRepresent(nodeA);
        Node<T> leaderB = findRepresent(nodeB);
        int sizeA = sizeMap.get(leaderA);
        int sizeB = sizeMap.get(leaderB);

        if(sizeA>sizeB){
            leaderMap.put(leaderB,leaderA);
            sizeMap.put(leaderA,sizeA+sizeB);
            sizeMap.remove(leaderB);
        }else{
            leaderMap.put(leaderA,leaderB);
            sizeMap.put(leaderB,sizeA+sizeB);
            sizeMap.remove(leaderA);
        }
    }

    public int unionSize(T t){
        return sizeMap.get(leaderMap.get(nodeMap.get(t)));
    }

    public int numOfSets(){
        return sizeMap.size();
    }
}
